Spiral
import torch
import matplotlib.pyplot as plt
import torch.nn as nn
from torchdiffeq import odeint_adjoint
from torchdiffeq import odeint
def spiral(t):
yx = torch.cos(t) * 2. * t
yy = torch.sin(t) * 2. * t
return torch.vstack((yx, yy))
x_train = torch.linspace(15,5,100)
y_train = spiral(x_train)
x_test = torch.linspace(5,1,30).reshape(-1)
y_test = spiral(x_test).T
plt.scatter(y_train[0], y_train[1], s=2, label='Train truth')
plt.scatter(y_test[:,0], y_test[:,1], s=2, label='Test truth')
plt.legend()
plt.show()
baseline = nn.Sequential(
nn.Linear(1, 100),
nn.ReLU(),
nn.Linear(100,100),
nn.ReLU(),
nn.Linear(100,100),
nn.ReLU(),
nn.Linear(100,100),
nn.ReLU(),
nn.Linear(100,100),
nn.ReLU(),
nn.Linear(100,2)
)
class Sine(nn.Module):
"""Sine activation with scaling.
Args:
w0 (float): Omega_0 parameter from SIREN paper.
"""
def __init__(self, w0=1.):
super().__init__()
self.w0 = w0
def forward(self, x):
return torch.sin(self.w0 * x)
baseline = nn.Sequential(
nn.Linear(1, 100),
Sine(),
nn.Linear(100,100),
Sine(),
nn.Linear(100,2)
)
def optimize_rough(net, x, y, steps):
optimizer = torch.optim.Adam(net.parameters())
criterion = nn.MSELoss()
losses = []
for i in range(steps):
optimizer.zero_grad()
y_pred = net(x)
loss = criterion(y_pred, y)
losses.append(loss.item())
loss.backward()
optimizer.step()
plt.scatter(torch.arange(len(losses)),losses)
plt.title('Training loss')
plt.show()
return net(x).detach()
y_pred = optimize_rough(baseline, x_train.reshape(-1,1), y_train.T, 10000)
y_pred_test = baseline(x_test.reshape(-1,1)).detach()
plt.plot(y_train[0], y_train[1], label='Train truth')
plt.plot(y_pred[:,0].detach(), y_pred[:,1], label='Baseline train')
plt.plot(y_test[:,0], y_test[:,1], label='Test truth')
plt.plot(y_pred_test[:,0], y_pred_test[:,1], label='Test model')
plt.legend()
plt.show()
class NeuralODE(nn.Module):
def __init__(self, phase_space_dim, projection_dim, hidden_layer_dim):
super().__init__()
self.delta_net = nn.Sequential(
nn.Linear(projection_dim + 1, hidden_layer_dim),
Sine(),
nn.Linear(hidden_layer_dim, hidden_layer_dim),
Sine(),
nn.Linear(hidden_layer_dim, hidden_layer_dim),
Sine(),
nn.Linear(hidden_layer_dim, projection_dim),
)
self.upprojector_net = nn.Linear(phase_space_dim + 1, projection_dim)
self.downprojector_net = nn.Linear(projection_dim, phase_space_dim)
def up(self, t, y):
concat = torch.hstack((t.reshape(1,1), y.reshape(1,-1)))
z = self.upprojector_net(concat)
return z.reshape(-1)
def down(self, z):
return self.downprojector_net(z.T).T
def forward(self, t, yt):
concat = torch.hstack((t.reshape(1,1), yt.reshape(1,-1)))
return self.delta_net(concat).reshape(-1)
def solve(self, y_start, t_start, t_stop, t_steps):
z_start = self.up(torch.tensor(t_start), y_start)
timepoints = torch.linspace(t_start,t_stop,t_steps)
z_pred = odeint(self, z_start, timepoints).T
return z_pred
nOde = NeuralODE(2,50,100)
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(nOde.parameters())
t_limit = 3
for i in range(1000):
print('Iteration', i)
optimizer.zero_grad()
z_pred = nOde.solve(y_train[:,0], 15., 5., 100) #odeint_adjoint(nOde, nOde.up(torch.tensor(0.), torch.randn(2)), torch.linspace(15,5,100)).T
y_pred = nOde.down(z_pred)
selection_indicator = (torch.rand(t_limit) > .2) #(torch.rand(y_train.shape[1]) > .2)
y_pred_selected = y_pred[:,:t_limit]#[:,selection_indicator] #(y_train.T[selection_indicator]).T
y_train_selected = y_train[:,:t_limit]#[:,selection_indicator]
print('selected', y_train_selected.shape)
loss = criterion(y_pred_selected, y_train_selected)
print('Loss:', loss)
plt.scatter(y_pred[0].detach()[:t_limit], y_pred[1].detach()[:t_limit], label='Euler neural')
plt.scatter(y_train[0][:t_limit], y_train[1][:t_limit], label='Truth')
plt.legend()
plt.show()
if loss.item() < 0.07:
t_limit += 1
print('Current t limit', t_limit)
loss.backward()
optimizer.step()
z_pred = nOde.solve(y_train[:,0], 15., 0., 200)
y_pred = nOde.down(z_pred)
y_pred.shape
x_full = torch.linspace(15.,0.,200)
y_full = spiral_easier(x_full)
y_full.shape
plt.scatter(y_full[0], y_full[1], label='Truth')
plt.scatter(y_pred[0].detach(), y_pred[1].detach(), label='nOde')
plt.legend()
torch.save(nOde, 'model1.pt')
Let's see if we can learn the sine wave function
x_train = torch.linspace(-10,10,100).reshape(-1,1)
y_train = torch.sin(x_train)
plt.scatter(x_train,y_train)
x_test = torch.linspace(10,20,50).reshape(-1,1)
y_test = torch.sin(x_test)
class NeuralODESin(nn.Module):
def __init__(self, phase_space_dim, projection_dim, hidden_layer_dim):
super().__init__()
self.delta_net = nn.Sequential(
nn.Linear(projection_dim + 1, hidden_layer_dim),
nn.ReLU(),
nn.Linear(hidden_layer_dim, hidden_layer_dim),
nn.ReLU(),
nn.Linear(hidden_layer_dim, hidden_layer_dim),
nn.ReLU(),
nn.Linear(hidden_layer_dim, projection_dim),
)
self.upprojector_net = nn.Linear(phase_space_dim + 1, projection_dim)
self.downprojector_net = nn.Linear(projection_dim, phase_space_dim)
def up(self, t, y):
concat = torch.hstack((t.reshape(1,1), y.reshape(1,-1)))
z = self.upprojector_net(concat)
return z.reshape(-1)
def down(self, z):
return self.downprojector_net(z)
def forward(self, t, yt):
concat = torch.hstack((t.reshape(1,1), yt.reshape(1,-1)))
return self.delta_net(concat).reshape(-1)
def solve(self, x_train, y_train):
y_start = y_train[0]
t_start = x_train[0]
z_start = self.up(torch.tensor(t_start), y_start)
z_pred = odeint(self, z_start, x_train.reshape(-1))
return z_pred
nde = NeuralODESin(1,50,100)
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(nde.parameters())
t_limit = 3
for i in range(10000):
optimizer.zero_grad()
t_start = torch.randint(high=max(1,t_limit-1),size=(1,))
x_train_batch = x_train[:t_limit] - x_train[t_start]
y_train_batch = y_train[:t_limit]
z_pred = nde.solve(x_train_batch, y_train_batch)
y_pred = nde.down(z_pred)
loss = criterion(y_pred, y_train_batch)
plt.plot(x_train_batch, y_train_batch, label='Truth')
plt.plot(x_train_batch, y_pred.detach(), label='Prediction')
plt.legend()
plt.show()
if loss.item() < 0.01:
t_limit += 1
print('Loss', loss.item())
print('Current t limit', t_limit)
loss.backward()
optimizer.step()
